import copy
import os
import warnings

import json5
import numpy as np
import trimesh

try:
    from part import Part, PART_DIR, GROUP_DIR,Global_Cache
except:
    from .part import Part, PART_DIR, GROUP_DIR,Global_Cache
import sympy as sym
from trimesh.transformations import unit_vector
# from dopy.transform import vector_rotation_matrix
from OCC.Extend.DataExchange import read_step_file
from OCC.Core.TopTools import TopTools_IndexedMapOfShape
from OCC.Extend.TopologyUtils import TopologyExplorer
from OCC.Core.TopoDS import topods_Vertex
from OCC.Core.BRep import BRep_Tool
try:
    from fillet import Rounding,Chamfer
    from transformations import identity_matrix,transform_points,concatenate_matrices,rotation_matrix,translation_matrix
except:
    from .fillet import Rounding, Chamfer
    from .transformations import identity_matrix, transform_points, concatenate_matrices, rotation_matrix, \
        translation_matrix


# from numba import jit


class SymPart(Part):
    def __init__(self, part_type, part_dir=PART_DIR, extra_args=None):
        # self.instance=Part(part_type,part_dir,extra_args,global_cache)
        self.parse_element_to_sym(part_type, part_dir)

        for name in self.elements:
            category = self.elements[name]['category']

            if category == '__box__':
                for k in ['x', 'y', 'z']:
                    k1 = name + '__' + k
                    if k1 not in self.args:
                        self.args[k1] = {}
                        self.syms[k1]=sym.symbols(k1)

                    if 'v' not in self.args[k1]:
                        self.args[k1]['v'] = 20
                        self.syms[k1]=20

            elif category in ['__cylinder__', '__cone__']:
                for k in ['r', 'h']:
                    k1 = name + '__' + k
                    if k1 not in self.args:
                        self.args[k1] = {}
                        self.syms[k1] = sym.symbols(k1)

                    if 'v' not in self.args[k1]:
                        self.args[k1]['v'] = 20
                        self.syms[k1] = 20
            elif category in ['__extrude__']:
                vertex_nums=int(self.args[name+"__vertex_num"]['v'])
                for i in range(vertex_nums):
                    self.syms[name+"__"+"p"+str(i)+"x"]=sym.symbols(name+"__"+"p"+str(i)+"x")
                    self.syms[name+"__"+'p'+str(i)+"y"]=sym.symbols(name+"__"+'p'+str(i)+"y")

            elif category =='__stp__':
                pass
        #
        self.parse_extra_args_sym(extra_args)
        #
        self.construct_sym()
        self.get_transform_sym()

    def parse_extra_args_sym(self, extra_args):
        if extra_args is None: return

        for ele in extra_args:
            name = ele['name']
            if name not in self.args:
                self.args[name] = {}
                self.syms[name] = sym.symbols(name)

            self.args[name].update(ele)


    def parse_element_to_sym(self, element_type, element_dir, prefix=''):
        self.element_type = element_type

        self.elements = {}
        self.args = {}
        self.points = {}
        self.trans_sym = {}
        self.syms={}

        if isinstance(element_type, dict):
            obj = element_type
        else:
            if prefix + element_type in Global_Cache.Cache:
                obj = Global_Cache.Cache[prefix + element_type]
            else:
                warnings.warn(prefix + element_type + ' not in Cache')
                with open(os.path.join(element_dir, element_type + ".json5"), "r", encoding="utf-8") as fp:
                    obj = json5.load(fp)

        assert 'elements' in obj
        if 'trans' not in obj: obj['trans'] = {}
        if 'points' not in obj: obj['points'] = {}
        if 'args' not in obj: obj['args'] = {}

        self.obj = obj

        for ele in obj['elements']:
            dic = {}
            for k in ['category', 'desc','address']:
                if k in ele: 
                    dic[k] = ele[k]

            dic['transform'] = trimesh.transformations.identity_matrix()
            dic['transform_sym']=sym.Matrix([[1, 0, 0,0], [0,1, 0, 0],[0,0,1,0],[0,0,0,1]])
            self.elements[ele['name']] = dic
        
        for ele in obj['args']:
            dic = {}
            for k in ele:
                dic[k] = ele[k]

            self.args[ele['name']] = dic
            self.syms[ele['name']]=sym.symbols(ele['name'])

        for ele in obj['points']:
            dic = {}
            for k in ['refer', 'offset']:
                if k in ele: 
                    dic[k] = ele[k]

            self.points[ele['name']] = dic

        for i, ele in enumerate(obj['trans']):
            name = ele.get('name', str(i))
            dic = {}
            for k in ['group', 'point2point', 'rotate', 'move', 'align']:
                if k in ele: 
                    dic[k] = copy.deepcopy(ele[k])
                    
            self.trans_sym[name] = dic

    def construct_sym(self):
        vs = {}
        symdict={}

        for k in self.args:
            vs[k] = self.args[k].get('v')
            symdict[k]=self.syms[k]

        for k in self.args:
            code = self.args[k].get('code')
            if code is not None:
                if isinstance(code,list):
                    for c in code:
                        exec(code.strip(), vs)
                        exec(code.strip(), symdict)
                else:
                    exec(code.strip(), vs)
                    exec(code.strip(),symdict)

        for k in self.args:
            self.args[k]['v'] = vs[k]
            self.syms[k]=symdict[k]

        for k1 in self.trans_sym:
            for k2 in self.trans_sym[k1]:
                if isinstance(self.trans_sym[k1][k2], (list, tuple)):
                    for i in range(len(self.trans_sym[k1][k2])):
                        if isinstance(self.trans_sym[k1][k2][i], str):
                            s = self.trans_sym[k1][k2][i]
                            if s.startswith('$'):
                                self.trans_sym[k1][k2][i] = self.syms[s[1:]]

    def get_point_sym(self, name):
        if isinstance(name, list):
            assert len(name) == 3
            return np.array(name, dtype=float)

        if '.' in name:
            ss = name.split('.')
            assert len(ss) == 2

            element = self.elements[ss[0]]

            position = ss[1]

            if element['category'] == '__box__':
                point_sym = sym.Matrix([[0, 0, 0,1]])
                vs_m=dict(
                    x=self.syms[ss[0] + '__' + 'x'],
                    y=self.syms[ss[0] + '__' + 'y'],
                    z=self.syms[ss[0] + '__' + 'z']

                )

                if '+x' in position:
                    point_sym[0]=vs_m['x']*0.5
                elif '-x' in position:
                    point_sym[0] = -vs_m['x'] * 0.5
                if '+y' in position:
                    point_sym[1] = vs_m['y'] * 0.5
                elif '-y' in position:
                    point_sym[1] = -vs_m['y'] * 0.5
                if '+z' in position:
                    point_sym[2] = vs_m['z'] * 0.5
                elif '-z' in position:
                    point_sym[2] = -vs_m['z'] * 0.5

                coor = element['transform_sym']*point_sym.T

                    # trimesh.transform_points([point], element['transform'])[0]

                return coor[0:3]
            elif element['category'] in ['__cylinder__', '__cone__']:
                point_sym = sym.Matrix([[0, 0, 0,1]])
                vs_m = dict(
                    r=self.syms[ss[0] + '__' + 'r'],
                    h=self.syms[ss[0] + '__' + 'h'],
                )

                if '+x' in position:
                    point_sym[0] = vs_m['r']
                elif '-x' in position:
                    point_sym[0] = -vs_m['r']
                if '+y' in position:
                    point_sym[1] = vs_m['r']
                elif '-y' in position:
                    point_sym[1] = -vs_m['r']
                if '+z' in position:
                    point_sym[2] = vs_m['h'] * 0.5
                elif '-z' in position:
                    point_sym[2] = -vs_m['h'] * 0.5

                coor = element['transform_sym']*point_sym.T

                return coor[0:3]

            elif element['category'] =='__stp__':
                m = read_step_file(element["address"])
                exp = TopologyExplorer(m)
                shape_map = TopTools_IndexedMapOfShape()
                for e in exp.vertices():
                    shape_map.Add(e)

                vertex = topods_Vertex(shape_map.FindKey(int(position)))
                # edge_face_map = TopTools_IndexedDataMapOfShapeListOfShape()
                aPnt = BRep_Tool.Pnt(vertex)
                coor=element['transform_sym']*np.array([aPnt.X(),aPnt.Y(),aPnt.Z(),1]).T
                return coor[0:3]
            elif element['category']=='__extrude__':
                position_index=int(position[2:])
                if '+' in position:
                    vertex=sym.Matrix([[self.syms[ss[0]+"__"+"p"+position[2:]+"x"],self.syms[ss[0]+"__"+"p"+position[2:]+"y"],self.syms[ss[0]+"__extrude"],1]])
                else:
                    vertex = sym.Matrix([[self.syms[ss[0]+"__"+"p"+position[2:]+"x"],self.syms[ss[0]+"__"+"p"+position[2:]+"y"], 0,1]])
                coor=element['transform_sym']*vertex.T
                return coor[0:3]
            else:
                raise NotImplementedError()
        else:
            assert name in self.points

            refer = self.points[name].get('refer')
            offset = self.points[name].get('offset')

            if refer is None:
                coor = sym.Matrix([[0, 0, 0]])
            else:
                coor = self.get_point_sym(refer)

            if offset is not None:
                coor += np.array(offset)

            return coor[0:3]
            
    def add_transform_sym(self, group, m):
        assert isinstance(group, list)
        
        for e in group:
            element = self.elements[e]
            element['transform_sym'] = m*element['transform_sym']


    def get_transform_sym(self):
        for ele in self.trans_sym.values():
            if 'align' in ele:
                a0 = [
                    self.get_point_sym(ele['align'][0] + '.+x'),
                    self.get_point_sym(ele['align'][0] + '.+y'),
                    self.get_point_sym(ele['align'][0] + '.+z')
                ]
                
                a1 = [
                    self.get_point_sym(ele['align'][1] + '.+x'),
                    self.get_point_sym(ele['align'][1] + '.+y'),
                    self.get_point_sym(ele['align'][1] + '.+z')
                ]

                a0 = a0 - self.get_point_sym(ele['align'][0] + '.o')
                a1 = a1 - self.get_point_sym(ele['align'][1] + '.o')
                
                m = vector_rotation_matrix(a0, a1)
                self.add_transform(ele['group'], m)

            if 'rotate' in ele:
                axis_map = {
                    "+x": Part.AXIS_X,
                    "+y": Part.AXIS_Y,
                    "+z": Part.AXIS_Z,
                }
                axis = axis_map[ele['rotate'][1].lower()]
                angle=sym.pi/180*ele['rotate'][0]
                sina = sym.sin(angle)
                cosa = sym.cos(angle)
                direction=axis[1]
                point=axis[0]
                direction = unit_vector(direction[:3])
                # rotation matrix around unit vector
                M = np.diag([cosa, cosa, cosa, 1.0])
                M[:3, :3] += np.outer(direction, direction) * (1.0 - cosa)

                direction = direction * sina
                M[:3, :3] += np.array([[0.0, -direction[2], direction[1]],
                                       [direction[2], 0.0, -direction[0]],
                                       [-direction[1], direction[0], 0.0]])

                # if point is specified, rotation is not around origin
                if point is not None:
                    point = np.array(point[:3], dtype=np.float64, copy=False)
                    M[:3, 3] = point - np.dot(M[:3, :3], point)

                m = sym.Matrix(M)
        
                self.add_transform_sym(ele['group'], m)

            if 'point2point' in ele:
                d1 = self.get_point_sym(ele['point2point'][0])
                d2 = self.get_point_sym(ele['point2point'][1])

                v = np.array(d2) - np.array(d1)
                # print(v)

                if len(ele['point2point']) >= 3:
                    keep = ele['point2point']
                    if 'x' not in keep:
                        v[0] = 0
                    if 'y' not in keep:
                        v[1] = 0    
                    if 'z' not in keep:
                        v[2] = 0 

                m = sym.Matrix([[1, 0, 0,v[0]], [0,1, 0, v[1]],[0,0,1,v[2]],[0,0,0,1]])

                self.add_transform_sym(ele['group'], m)

            if 'move' in ele:
                v=ele['move']
                m = sym.Matrix([[1, 0, 0,v[0]], [0,1, 0, v[1]],[0,0,1,v[2]],[0,0,0,1]])
                
                self.add_transform_sym(ele['group'], m)




if __name__ == '__main__':
    os.chdir(os.path.dirname(__file__))
    # print(len(Global_Cache.Cache))



    x = SymPart("new_ddfa")
    test=x.get_point_sym('A.+p2')
    print(test)
    # test=x.get_point_sym('A.-z')
    # print(test)
    # # print(test)
    # L1,L2,SIZE,H,ARC=sym.symbols('L1,L2,SIZE,H,ARC')
    #
    # f=sym.lambdify((L1,L2,SIZE,H,ARC), test, 'numpy')
    # # g = jit(f)
    #
    # start=time.time()
    # tes= [sym.N(e.subs({SIZE: 20, L1: 80, L2: 50, H: 25, ARC: 105})) for e in x.get_point_sym('A.-z')]
    # dt=time.time()-start
    # print(dt)
    # # print(test)
    #
    # start=time.time()
    # print(f(80,50,20,25,105))
    # dt=time.time()-start
    # print(dt)
    # print(g(20,80,50,25,105))


    # test=[sym.N(e.subs({SIZE:20,L1:80,L2:50,H:25,ARC:105})) for e in x.get_point_sym('A.-z')]
    # print(test)
    # print(x.get_point_sym('A.-z'))
    # x1=Part('cY1')
    # print(x1.get_point('A.-z'))
    # print(x.syms)
    # Global_Cache.Cache
    # x.show()
    # sys.exit()
    #
    # # from shapely.geometry import Polygon
    # # polygon = Polygon([[0, 0], [1, 0], [1, 3]])
    # # # polygon = trimesh.path.polygons.paths_to_polygons([[0, 0], [0, 1], [1, 1]])
    # # m = trimesh.creation.extrude_polygon(polygon, 0.5)
    # # m.show()
    # # sys.exit()
    # x = Part("cY", extra_args=[{'name': "ARC", 'v': 60}])
    # # x.show()
    # x.show([x.get_point("A1")])
    # sys.exit()
    #
    # x = Part("cPad")
    #
    # vs = x.get_parameters()
    # pprint(vs)
    # x.show()
    #
    # # m = x.get_mesh()
    # sys.exit()
    # # pprint(vs)
    #
    # from CATDocument import CATDocument
    # cat = CATDocument("L1.CATPart")
    # cat.open()
    #
    # for k in vs:
    #     try:
    #         cat.write_parameter_new(k, vs[k])
    #     except:
    #         pass
    #
    # cat.update()


    # cat.close()
    
